import uuid
import threading
from typing import Optional, Any

from config import EMBEDDING_DIM_SIZE

import faiss
import numpy as np

class BidirectionalIdMapper:
    def __init__(self):
        self._lock = threading.Lock()
        self._uuid_to_int: dict[uuid.UUID, int] = {}
        self._int_to_uuid: dict[int, uuid.UUID] = {}
        self._counter: int = 0

    def get_or_create_int_id(self, uuid_id: uuid.UUID) -> int:
        with self._lock:
            if uuid_id not in self._uuid_to_int:
                int_id = self._counter
                self._uuid_to_int[uuid_id] = int_id
                self._int_to_uuid[int_id] = uuid_id
                self._counter += 1
                return int_id
            return self._uuid_to_int[uuid_id]

    def get_uuid_id(self, int_id: int) -> Optional[uuid.UUID]:
        with self._lock:
            return self._int_to_uuid.get(int_id)

    def remove_mappings(self, uuid_ids: list[uuid.UUID]):
        with self._lock:
            for uuid_id in uuid_ids:
                if uuid_id in self._uuid_to_int:
                    int_id = self._uuid_to_int.pop(uuid_id)
                    self._int_to_uuid.pop(int_id, None)


class FaissManager:
    def __init__(self, dim: int = EMBEDDING_DIM_SIZE):
        self._lock = threading.Lock()
        self._mapper = BidirectionalIdMapper()
        self._faiss_index = faiss.IndexIDMap2(faiss.IndexFlatIP(dim))

    def add(self, vectors_data: list[dict[str, Any]]):
        if not vectors_data:
            return

        int_ids = [self._mapper.get_or_create_int_id(item['id']) for item in vectors_data]
        embeddings = [item['embedding'] for item in vectors_data]

        np_embeddings = np.array(embeddings, dtype=np.float32)
        faiss.normalize_L2(np_embeddings)

        np_ids = np.array(int_ids, dtype=np.int64)

        with self._lock:
            self._faiss_index.add_with_ids(np_embeddings, np_ids)

    def remove(self, ids: list[uuid.UUID]):
        if not ids:
            return

        int_ids_to_remove = [self._mapper.get_or_create_int_id(uid) for uid in ids]
        np_ids_to_remove = np.array(int_ids_to_remove, dtype=np.int64)

        with self._lock:
            self._faiss_index.remove_ids(np_ids_to_remove)
            self._mapper.remove_mappings(ids)

    def search_top_k(self, query_vector: list[float], id_selector: list[uuid.UUID], k: int) -> list[tuple[uuid.UUID, float]]:
        if not id_selector:
            return []

        np_query = np.array([query_vector], dtype=np.float32)
        faiss.normalize_L2(np_query)

        int_id_selector = [self._mapper.get_or_create_int_id(uid) for uid in id_selector]
        if not int_id_selector:
            return []

        np_id_selector = np.array(int_id_selector, dtype=np.int64)

        selector = faiss.IDSelectorBatch(np_id_selector)

        search_params = faiss.SearchParameters()
        search_params.sel = selector

        with self._lock:
            # k 参数现在是动态的
            similarities, indices = self._faiss_index.search(np_query, k=k, params=search_params)

        results = []
        for sim, idx in zip(similarities[0], indices[0]):
            if idx != -1:
                found_uuid = self._mapper.get_uuid_id(idx)
                if found_uuid:
                    results.append((found_uuid, float(sim)))
        return results

    def search(self, query_vector: list[float], id_selector: list[uuid.UUID]) -> Optional[tuple[uuid.UUID, float]]:
        results = self.search_top_k(query_vector, id_selector, k=1)
        return results[0] if results else None

    @property
    def total_vectors(self) -> int:
        return self._faiss_index.ntotal